"""
带噪声的识别、位姿解算测试功能
"""

import numpy as np
from match.binad import XuPyramid, ChenPyramid, load_for_match, fit_ellipse
import os
import tqdm
from scipy.spatial.transform import Rotation
import argparse
import yaml


class GaussianCrater:
    def __init__(self, *args):
        self.mu = []
        self.sigma = []
        for arg in args:
            self.mu.append(np.mean(arg))
            self.sigma.append(np.std(arg))

    def __call__(self):
        return [np.random.normal(mu, sigma) for mu, sigma in zip(self.mu, self.sigma)]


def arg_parse():
    arg_parser = argparse.ArgumentParser()
    arg_parser.add_argument(
        "-c",
        "--collect_dir",
        type=str,
        help="The directory of the collected data",
    )
    arg_parser.add_argument(
        "-o",
        "--catalog_dir",
        type=str,
        default="/disk527/sdb1/a804_cbf/catalog",
        help="The directory of the catalog",
    )
    arg_parser.add_argument(
        "--center",
        type=float,
        default=0.5,
    )
    arg_parser.add_argument(
        "--edge",
        type=float,
        default=0.5,
    )
    arg_parser.add_argument(
        "--add",
        type=float,
        default=0.1,
    )
    arg_parser.add_argument(
        "--rm",
        type=float,
        default=0.05,
    )
    arg_parser.add_argument(
        "--sigma",
        type=float,
        default=0.3,
    )
    arg_parser.add_argument(
        "-m",
        "--method",
        choices=["xuliheng", "chenbofei"],
        help="The method to match the crater",
    )
    return arg_parser.parse_args()


def noiser(points, center: float = 0.5, edge: float = 0.5, sigma: float = 1):
    """处理中心偏离和边缘偏离"""
    if np.random.rand() < center:
        points = points + np.random.normal(0, sigma)
    if np.random.rand() < edge:
        points = points + np.random.normal(0, sigma, size=points.shape)
    return points


def ellipse_params(idx, elli):
    """
    Arguments:
        idx (list): 陨石坑的id
        elli (np.ndarray): 陨石坑的椭圆参数 N x 6
    """
    # F归一化
    elli = elli / elli[:, 5, None]
    cnt_x = (elli[:, 1] * elli[:, 4] - 2 * elli[:, 2] * elli[:, 3]) / (
        4 * elli[:, 0] * elli[:, 2] - elli[:, 1] ** 2
    )
    cnt_y = (elli[:, 1] * elli[:, 3] - 2 * elli[:, 0] * elli[:, 4]) / (
        4 * elli[:, 0] * elli[:, 2] - elli[:, 1] ** 2
    )
    theta = -0.5 * np.arctan(2 * elli[:, 1] / (elli[:, 0] - elli[:, 2]))
    nC = np.array(
        [
            [elli[:, 0], elli[:, 1] / 2, elli[:, 3] / 2],
            [elli[:, 1] / 2, elli[:, 2], elli[:, 4] / 2],
            [elli[:, 3] / 2, elli[:, 4] / 2, elli[:, 5]],
        ]
    ).T
    ab = np.sqrt(
        (
            nC[:, 2, 2, None]
            - (nC[:, 2, None, :2] @ nC[:, :2, :2] @ nC[:, :2, 2, None]).squeeze(-1)
        )
        / np.linalg.eigvals(nC[:, :2, :2])
    )
    ind = ~(
        np.isnan(cnt_x)
        | np.isnan(cnt_y)
        | np.isnan(ab[:, 0])
        | np.isnan(ab[:, 1])
        | np.isnan(theta)
    )
    idx = np.array(idx)[ind]
    return idx, cnt_x[ind], cnt_y[ind], ab[ind, 0], ab[ind, 1], theta[ind]


def params_ellipse(gen: GaussianCrater):
    """
    以椭圆六个参数(x^2, xy, y^2, x, y, 1)为返回值
    """
    while True:
        cnt_x, cnt_y, a, b, theta = gen()
        elli = [
            a * np.cos(theta) ** 2 + b * np.sin(theta) ** 2,
            (b - a) * np.sin(theta) * np.cos(theta),
            a * np.sin(theta) ** 2 + b * np.cos(theta) ** 2,
            -2 * a * cnt_x * np.cos(theta) - b * cnt_y * np.sin(theta),
            -2 * b * cnt_y * np.cos(theta) - a * cnt_x * np.sin(theta),
            a * cnt_x**2 + b * cnt_y**2 - a * b,
        ]
        if np.isnan(elli).any():
            continue
        return elli


def misdetect(idx: list, params: list, prob_add: float = 0.1, prob_del: float = 0.05):
    """
    以当前图像平面上的陨石坑中心、长轴、短轴、旋转角的分布为独立随机分布，以概率prob产生误检测
    Arguments:
        prob_add (float): 误增加概率
        prob_del (float): 误删除概率
    """
    flag = [False, False]
    ids = []
    ellis = []
    for id, param in zip(idx, params):
        if np.random.rand() < prob_del:
            flag[0] = True
        if np.random.rand() < prob_add:
            param_false = params_ellipse(crater_gen)
            flag[1] = True
        if flag[0] and flag[1]:
            ellis.append(param_false)
            ids.append(id)
        elif not flag[0] and flag[1]:
            ellis.append(param_false)
            ellis.append(param)
            ids.append("false")
            ids.append(id)
        elif not flag[0] and not flag[1]:
            ellis.append(param)
            ids.append(id)
        flag = [False, False]
    return ids, ellis


if __name__ == "__main__":
    args = arg_parse()
    if args.method == "xuliheng":
        method = XuPyramid
    elif args.method == "chenbofei":
        method = ChenPyramid
    xu_match = method(args.catalog_dir)
    params = xu_match.catalog[:, 3:]
    _, *crater_params = ellipse_params(np.arange(params.shape[0]), params)
    crater_gen = GaussianCrater(*crater_params)
    points_dir = os.path.join(args.collect_dir, "points_id")
    Acc = []
    Recall = []
    err_dist = []
    err_rot = []
    it = 0
    test_iter = 0
    output_file = "_".join(
        f"{k}_{v}" for k, v in args.__dict__.items() if "dir" not in k
    )
    with open(os.path.join(args.collect_dir, "config.yaml"), "r") as f:
        config = yaml.load(f.read(), Loader=yaml.FullLoader)
    K = np.array(config["camera"]["P"], dtype=np.float32).reshape(3, 4)[:3, :3]
    with open(output_file + ".csv", "w") as writer:
        print(
            "time/s,x/m,y/m,z/m,q0,q1,q2,q3,accuracy/%,recall/%,success/%,dist_x/m,dist_y/m,dist_z/m,rot_x/deg,rot_y/deg,rot_z/deg",
            file=writer,
        )
        with open(f"{args.collect_dir}/pose.csv", "r") as f:
            f.readline()  # 去除标题
            file_names = sorted(
                os.listdir(points_dir), key=lambda x: float(x.split(".txt")[0])
            )
            for time in tqdm.tqdm(
                map(lambda x: x.split(".txt")[0], file_names), total=len(file_names)
            ):
                if float(time) < 100:
                    continue
                it += 1
                for line in f:
                    # 这个判断语句有点问题，如果把阈值改小到5e-4，就会在某个点处产生错误，然后因为不能更新flag，导致后续的所有点都会出错
                    if abs(float(line.split(",")[0]) - float(time)) < 2e-4:
                        pose = [float(x) for x in line.strip().split(",")[1:]]
                        break
                    elif float(line.split(",")[0]) > float(time):
                        pose = [float(x) for x in line.strip().split(",")[1:]]
                        break
                ellipsis = []
                gt_id = []
                for idx, points in load_for_match(
                    os.path.join(points_dir, time + ".txt"),
                    args.center,
                    args.edge,
                    args.sigma,
                ):
                    if len(points) < 5:
                        continue
                    elli, err = fit_ellipse(np.array(points).T, return_error=True)
                    if err < 1e-3:
                        ellipsis.append(elli)
                        gt_id.append(idx)
                if len(ellipsis) < 4:
                    continue
                # 产生误检测
                gt_id, ellipsis = misdetect(gt_id, ellipsis, args.add, args.rm)
                # 将陨石坑顺序打乱
                rand_idx = np.random.permutation(len(gt_id))
                gt_id = np.array(gt_id)[rand_idx]
                ellipsis = np.array(ellipsis)[rand_idx]
                for (ijkl, name), (R, T), error in xu_match(
                    ellipsis, K, factor=0.3, pair_th=1000, triple_th=1000, reprj_th=0.2
                ):
                    if error < 0.15:
                        test_iter += 1
                        acc = 0
                        for i, j in zip(map(lambda i: gt_id[i], ijkl), name):
                            if i == j:
                                acc += 1
                        Acc.append(acc / len(name))
                        Recall.append(len(name) / len(gt_id))
                        ##### Proof for pose calculation #####
                        # 计算相机在世界坐标系下的位置，已知从世界坐标系转到相机坐标系的旋转矩阵R和平移向量T，相机中心为C
                        # 则有：T = - R @ C => C = - R.T @ T
                        # 计算相机坐标系相对于世界坐标系的三轴转角（用四元数表示）。OpenCV建系认为x轴沿图像宽，y轴沿图像高，且
                        # 对于OpenCV图像上一点齐次点u，它与真实图像系齐次点v的关系为：u = S @ v
                        # 而真实图像系齐次点与相机坐标系三维点X的投影关系为：v = K @ X
                        # 相机坐标系三维点与世界坐标系三维点的关系为：X = R' @ X' + T'
                        # 故有：u = S @ K @ R' @ X' + S @ K @ T'
                        # 特别地，由于K是上三角矩阵且右上角元素相等，因此 S @ K = K @ S
                        # 故有：u = K @ S @ R' @ X' + K @ S @ T'
                        # 因此OpenCV中利用solvePnP解算的旋转矩阵R和平移向量T分别等于：
                        # R = S @ R', T = S @ T'
                        # 下面证明当前计算位置和姿态的过程是正确的：
                        # 求解真实空间的相机位置，即C = - R.T @ T = - (R'.T @ S.T) @ (S @ T') = - R' @ T' (由于S是初等变换矩阵)
                        # 求解真实空间的相机姿态，即R.T = R'.T @ S.T = R'.T @ S
                        # -1的由来：旋转矩阵的行列式只能为1，不能为-1，而S的行列式显然为-1并不是旋转矩阵，因此需要将其校正为行列式为1的旋转矩阵
                        # 由于S正交、对称，因此将其取反即可保证行列式为1，故：
                        # R =  R'.T @ (-S) = -R'.T @ S
                        ##### End of proof #####
                        pose_esti = -R.T @ T
                        # 取负才是行列式为1的旋转矩阵
                        S = -np.array([[0, 1, 0], [1, 0, 0], [0, 0, 1]])
                        ori = Rotation.from_matrix(R.T @ S).as_quat()
                        err_dist.append(
                            (pose_esti.squeeze() - np.array(pose[:3])).tolist()
                        )
                        err_rot.append(
                            (
                                Rotation.from_quat(ori)
                                * Rotation.from_quat(pose[3:]).inv()
                            )
                            .as_rotvec(degrees=True)
                            .tolist()
                        )
                        print(
                            time,
                            ",".join(map(str, pose)),
                            str(acc / len(name)),
                            str(len(name) / len(gt_id)),
                            ",".join(map(str, err_dist[-1])),
                            ",".join(map(str, err_rot[-1])),
                            sep=",",
                            file=writer,
                            flush=True,
                        )
                        break
                if it % 100 == 0 and len(err_dist) > 0:
                    print(
                        f"Current accuracy is {np.mean(Acc):4f}, recall is {np.mean(Recall):4f}, success rate is {(test_iter/it):3f}"
                    )
                    temp_dist = np.sqrt(
                        np.mean(np.sum(np.array(err_dist) ** 2, axis=1))
                    )
                    temp_rot = np.sqrt(np.mean(np.sum(np.array(err_rot) ** 2, axis=1)))
                    print(
                        f"Current distance error is {temp_dist:4f}, rotation error is {temp_rot:4f}",
                    )

    print(f"Accuracy is {np.mean(Acc)}, recall is {np.mean(Recall)}")
    print(
        f"Distance error is {np.sqrt(np.mean(np.sum(np.array(err_dist)**2,axis=1)))}, rotation error is {np.sqrt(np.mean(np.sum(np.array(err_rot)**2,axis=1)))}"
    )
    print(f"Test success rate is {test_iter/it}")
